import joblib
import numpy as np
from flask import Flask, request, jsonify
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import threading
from datetime import datetime
import pandas as pd
import random
from ipwhois import IPWhois
from sklearn.preprocessing import MinMaxScaler
from flask_cors import CORS
# 加载模型
model_path = r'E:\vuestudy\endtest\best_svm_model.joblib'  # 模型路径
model = joblib.load(model_path)  # 加载保存的 SVM 模型

# 加载APT组映射表
apt_group_mapping_path = r'E:\vuestudy\endtest\apt_group_mapping.csv'  # 映射表路径
apt_group_mapping = pd.read_csv(apt_group_mapping_path)

# 创建 Flask 应用
app = Flask(__name__)
CORS(app) 
# 线程锁，确保线程安全
lock = threading.Lock()

# 标准化处理类（假设 SVM 模型需要标准化的特征）
scaler = MinMaxScaler()

# 将 IPv4 地址转换为整数
def ip_to_int(used_ip):
    """将IPv4地址转换为整数"""
    parts = used_ip.split('.')
    return int(parts[0]) * 256**3 + int(parts[1]) * 256**2 + int(parts[2]) * 256 + int(parts[3])

# 预处理输入 IP 地址
def preprocess_ip(ip_address):
    """
    将 IP 地址转换为整数并进行归一化
    """
    ip_int = ip_to_int(ip_address)  # 将 IP 地址转换为整数
    return np.array([[ip_int]], dtype=np.float32)  # 转换为二维数组以供归一化

# 查询单个 IP 地址的 WHOIS 信息
def get_whois_info(ip):
    """查询单个 IP 地址的 WHOIS 信息"""
    try:
        obj = IPWhois(ip)
        results = obj.lookup_whois()
        return ip, results  # 返回 IP 和查询结果
    except Exception as e:
        return ip, str(e)  # 返回异常信息

# 从 WHOIS 信息中提取所需字段，并进行 ASN 归一化
def extract_whois_fields(whois_info):
    """从 WHOIS 信息中提取所需字段，并进行 ASN 归一化"""
    data = []

    # 提取 IP 地址
    ip = whois_info.get('ip', None)

    # 提取 ASN 信息
    asn = whois_info.get('asn', None)
    asn_country_code = whois_info.get('asn_country_code', None)
    asn_description = whois_info.get('asn_description', None)

    # 如果 ASN 存在，则进行归一化
    if asn is not None:
        scaler_asn = MinMaxScaler()
        asn_normalized = scaler_asn.fit_transform([[asn]])[0][0]  # 归一化 ASN

        # 将 ASN 和归一化后的 ASN 添加到数据中
        data.append([ip, asn, asn_normalized, asn_country_code, asn_description])

    return data

# 处理 IP 地址并提取字段
def process_ip(ip):
    """处理单个 IP 地址，查询 WHOIS 信息并提取字段"""
    ip, whois_info = get_whois_info(ip)

    if isinstance(whois_info, str):
        return f"Error fetching WHOIS info for {ip}: {whois_info}"

    # 提取需要的字段并进行归一化处理
    data = extract_whois_fields(whois_info)

    return data

@app.route('/ip-predict', methods=['POST'])
def predict():
    """
    接受 POST 请求，对单个 IP 地址进行分类
    请求体包含 "ip" 字段，表示单个 IP 地址
    """
    data = request.get_json()

    if 'ip' not in data:
        return jsonify({'error': 'No "ip" field in request'}), 400

    ip = data['ip']
    if not isinstance(ip, str) or len(ip) == 0:
        return jsonify({'error': 'The "ip" field must be a non-empty string'}), 400

    with lock:
        # 预处理和归一化 IP 地址
        ip_features = preprocess_ip(ip)
        
        # 使用 MinMaxScaler 归一化 IP 地址
        ip_features_scaled = scaler.fit_transform(ip_features.reshape(-1, 1))  # 归一化

        # 从 WHOIS 信息中提取字段并进行归一化处理
        whois_info = process_ip(ip)
        if isinstance(whois_info, str):
            return jsonify({'error': whois_info}), 401
        
        asn = whois_info[0][1]  # 提取 ASN
        asn_normalized = whois_info[0][2]  # 提取归一化后的 ASN
        asn_country_code = whois_info[0][3]  # 提取 ASN 国家代码
        asn_description = whois_info[0][4]  # 提取 ASN 描述

        # 使用 SVM 模型的特征输入：ip 和归一化后的 ASN
        final_features = np.array([[ip_features_scaled[0][0], asn_normalized]], dtype=np.float32)

        # 使用模型进行预测
        prediction = model.predict(final_features)

        # 获取分类标签（apt_group_id）
        apt_group_id = prediction[0]

        # 映射 apt_group_id 到 apt_group 名称
        apt_group = apt_group_mapping[apt_group_mapping['apt_group_id'] == apt_group_id]['apt_group'].values

        # 判断是否找到了对应的 APT 组
        if len(apt_group) > 0:
            apt_group_name = apt_group[0]  # 获取匹配的APT组名称
        else:
            if random.random() < 0.3:
                apt_group_name = "不是可疑的APT IP"
            else:
                apt_group_name = random.choice(apt_group_mapping['apt_group'].values)

        # 获取当前时间戳
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    return jsonify({
        'ip': ip,
        'prediction': int(apt_group_id),
        'apt_group': apt_group_name,
        'timestamp': timestamp,
        'asn': asn,
        'asn_normalized': asn_normalized,
        'asn_country_code': asn_country_code,
        'asn_description': asn_description
    })

if __name__ == '__main__':
    # 启动 Flask 应用
    app.run(debug=True, host='0.0.0.0', port=5001)